import os
import torch
import timm
import torchvision.transforms as transforms
from flask import Flask, render_template, request, jsonify
from PIL import Image
import io
import numpy as np
from flask_cors import CORS
import torch.nn as nn

# 加载训练好的EfficientNet-B4模型
def load_model(model_path, num_classes):
    # 初始化EfficientNet-B4模型
    model = timm.create_model('tf_efficientnet_b4', pretrained=False)
    
    # 修改分类层
    model.classifier = nn.Sequential(
        nn.Dropout(0.5),  # 添加Dropout正则化
        nn.Linear(model.classifier.in_features, num_classes)
    )

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件 `{model_path}` 不存在，请检查路径是否正确！")

    try:
        model.load_state_dict(torch.load(model_path))
        print("模型权重加载成功！")
    except Exception as e:
        print(f"模型权重加载失败: {e}")
        raise

    model.eval()
    return model

# 初始化 Flask 应用
app = Flask(__name__, template_folder='templates')  # 指定模板文件夹
CORS(app)

# 模型路径和类别
model_path = 'D:\\深度学习\\期末设计\\best_efficientnet_b4_model.pth'  # 确保这是s.py训练后保存的模型路径
num_classes = 36  # 确保类别数与s.py中的相同
model = load_model(model_path, num_classes)

# 类别映射（确保与s.py中的类别一致）
mushroom_classes = [
    '口蘑', '奶浆菌', '姬松茸', '干巴菌', '平菇',
    '杏鲍菇', '松茸', '松露', '榆黄蘑', '榛蘑',
    '牛肝菌', '猪肚菌', '猴头菇', '白参菌', '白玉菇',
    '白葱牛肝菌', '竹荪', '红菇', '羊肚菌', '老人头菌',
    '茶树菇', '草菇', '虎掌菌', '虫草花', '蟹味菇',
    '谷熟菌', '金耳', '金针菇', '银耳', '青头菌',
    '香菇', '鸡枞菌', '鸡油菌', '鸡腿菇', '鹿茸菇', '黑木耳'
]

# 设备设置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前设备: {device}")
model = model.to(device)

# 定义与训练一致的基础预处理（与train.py中的transform一致）
transform = transforms.Compose([
    transforms.Resize((380, 380)),  # EfficientNet-B4推荐的输入尺寸
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(45),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
    transforms.RandomPerspective(distortion_scale=0.2, p=0.5),  # 添加透视变换
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 主页路由
@app.route('/')
def home():
    return render_template('index.html')  # 渲染原有的前端页面

# 前端页面路由
@app.route('/upload')
def upload():
    return render_template('upload.html')  # 渲染新的前端页面

# 图片上传和预测接口
@app.route('/predict', methods=['POST'])
def predict():
    try:
        if 'image' not in request.files:
            return jsonify({'error': '未找到图片文件'}), 400

        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '未选择图片文件'})

        img_bytes = file.read()
        img = Image.open(io.BytesIO(img_bytes)).convert('RGB')

        # 应用与训练一致的预处理
        img = transform(img)

        with torch.no_grad():
            output = model(img.unsqueeze(0).to(device))
            probabilities = torch.nn.functional.softmax(output, dim=1)
            predicted_class_idx = torch.argmax(probabilities).item()
            confidence = probabilities[0][predicted_class_idx].item()

        predicted_class = mushroom_classes[predicted_class_idx]

        print(f"预测结果: {predicted_class}, 置信度: {confidence:.4f}")
        
        # 返回预测结果和置信度
        return jsonify({
            'prediction': predicted_class,
            'confidence': confidence
        })

    except Exception as e:
        print(f"预测错误: {e}")
        return jsonify({'error': str(e)}), 500

# 批量预测接口
@app.route('/batch_predict', methods=['POST'])
def batch_predict():
    try:
        files = request.files.getlist('images')  # 获取多个文件
        if not files:
            return jsonify({'error': '未找到图片文件'}), 400

        results = []
        for file in files:
            if file.filename == '':
                continue

            img_bytes = file.read()
            img = Image.open(io.BytesIO(img_bytes)).convert('RGB')

            img = transform(img)

            with torch.no_grad():
                output = model(img.unsqueeze(0).to(device))
                probabilities = torch.nn.functional.softmax(output, dim=1)
                predicted_class_idx = torch.argmax(probabilities).item()
                confidence = probabilities[0][predicted_class_idx].item()

            predicted_class = mushroom_classes[predicted_class_idx]

            results.append({
                'filename': file.filename,
                'prediction': predicted_class,
                'confidence': confidence
            })

        return jsonify({'results': results})

    except Exception as e:
        print(f"批量预测错误: {e}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)  # 使用5000端口